# m2_pytorch_mps_demo.py

import torch
import time

def check_mps_availability():
    """检查MPS可用性"""
    print("=== MPS可用性检查 ===")
    if torch.backends.mps.is_available():
        print("MPS可用，可以使用GPU加速")
        device = torch.device("mps")
    else:
        print("MPS不可用，将使用CPU进行计算")
        device = torch.device("cpu")
    print(f"使用的设备: {device}")
    print()
    return device

def demonstrate_mps_computation(device):
    """演示MPS计算"""
    print("=== MPS计算演示 ===")
    
    # 创建张量
    x_cpu = torch.tensor([1., 2., 3.])
    print(f"CPU张量: {x_cpu} (设备: {x_cpu.device})")
    
    if device.type == "mps":
        # 将张量移到MPS设备
        x_mps = x_cpu.to(device)
        print(f"MPS张量: {x_mps} (设备: {x_mps.device})")
        
        # 在MPS上执行计算
        y_mps = x_mps * 2
        print(f"MPS计算结果: {y_mps} (设备: {y_mps.device})")
        
        # 将结果移回CPU
        y_cpu = y_mps.to("cpu")
        print(f"移回CPU的结果: {y_cpu} (设备: {y_cpu.device})")
    else:
        print("由于MPS不可用，只能在CPU上进行计算")
        y_cpu = x_cpu * 2
        print(f"CPU计算结果: {y_cpu} (设备: {y_cpu.device})")
    print()

def demonstrate_performance_difference(device):
    """演示MPS和CPU性能差异"""
    print("=== MPS与CPU性能差异演示 ===")
    
    if device.type == "mps":
        # 创建大尺寸张量
        size = 2000
        x_cpu = torch.randn(size, size)
        y_cpu = torch.randn(size, size)
        
        # CPU计算
        start_time = time.time()
        z_cpu = torch.mm(x_cpu, y_cpu)
        cpu_time = time.time() - start_time
        print(f"CPU计算时间: {cpu_time:.4f}秒")
        
        # MPS计算
        x_mps = x_cpu.to(device)
        y_mps = y_cpu.to(device)
        
        # 预热MPS
        _ = torch.mm(x_mps, y_mps)
        torch.mps.synchronize()
        
        start_time = time.time()
        z_mps = torch.mm(x_mps, y_mps)
        torch.mps.synchronize()  # 等待MPS计算完成
        mps_time = time.time() - start_time
        print(f"MPS计算时间: {mps_time:.4f}秒")
        
        print(f"加速比: {cpu_time/mps_time:.2f}x")
    else:
        print("MPS不可用，无法演示GPU加速效果")
    print()

def explain_mps_info():
    """解释MPS相关信息"""
    print("=== MPS相关信息 ===")
    print("Metal Performance Shaders (MPS) 是苹果为Metal API提供的高性能计算框架")
    print("它可以利用M1/M2芯片的GPU进行加速计算")
    print()
    print("主要特点:")
    print("  - 专为苹果芯片优化")
    print("  - 支持大部分PyTorch操作")
    print("  - 与CUDA类似的设备管理方式")
    print("  - 自动内存管理")
    print()

def explain_mps_installation():
    """解释MPS安装选项"""
    print("=== MPS安装说明 ===")
    print("对于Apple Silicon Mac (M1/M2)，安装PyTorch只需要:")
    print("  pip install torch torchvision torchaudio")
    print()
    print("注意事项:")
    print("  - 需要macOS 12.3或更高版本")
    print("  - 需要PyTorch 1.12或更高版本")
    print("  - 确保使用Python 3.8或更高版本")
    print()

def explain_memory_management(device):
    """解释MPS内存管理"""
    print("=== MPS内存管理 ===")
    if device.type == "mps":
        print("MPS内存管理特点:")
        print("  - PyTorch自动管理MPS内存分配")
        print("  - 张量创建时分配MPS内存")
        print("  - 张量销毁时释放MPS内存")
        print()
        print("内存管理函数:")
        print("  - torch.mps.empty_cache() - 释放未使用的缓存")
        print("  - torch.mps.memory_allocated() - 已分配内存")
        print("  - torch.mps.driver_allocated_memory() - 驱动分配内存")
    else:
        print("MPS不可用，无法演示MPS内存管理")
    print()

def main():
    """主函数"""
    print("PyTorch MPS支持演示 (Apple M1/M2)")
    print("=" * 50)
    
    # 检查MPS可用性
    device = check_mps_availability()
    
    # 演示MPS计算
    demonstrate_mps_computation(device)
    
    # 演示性能差异
    demonstrate_performance_difference(device)
    
    # 解释MPS相关信息
    explain_mps_info()
    
    # 解释安装选项
    explain_mps_installation()
    
    # 解释内存管理
    explain_memory_management(device)
    
    print("=" * 50)
    print("程序执行完成!")

if __name__ == '__main__':
    main()
